# imports
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# create convolutional neural networks
class CNN(nn.Module):
    def __init__(self, in_channels, num_classes):
        super(CNN, self).__init__()
        self.cnn1 = nn.Conv2d(in_channels=in_channels, out_channels=8, kernel_size=3, stride=1, padding=1)
        self.cnn2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(in_features=16 * 7 * 7, out_features=num_classes)
    def forward(self, x):
        # (1x28x28)
        x = F.relu(self.cnn1(x)) #(8x28x28)
        x = self.pool(x) #(8x14x14)
        x = F.relu(self.cnn2(x)) #(16x14x14)
        x = self.pool(x) #(16x7x7)
        return self.fc(x.view(x.shape[0], -1))

# train
def train(model, loss_fn, optimizer, dataloader, epochs):
    model.train()
    print("=" * 10 + "training" + "=" * 10)
    for epoch in range(epochs):
        correct_samples = 0
        total_samples = 0
        total_loss = 0

        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)
            # data = data.reshape(data.shape[0], -1)
            result = model(data)
            loss = loss_fn(result, target)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            total_samples += target.shape[0]
            correct_samples += target.eq(result.argmax(-1)).sum().item()
        print(f"epoch: {epoch}. loss: {total_loss/ total_samples}, accuracy: {correct_samples / total_samples * 100:.2f}% ")


def test(model, loss_fn, dataloader):
    model.eval()
    print("=" * 10 + "testing" + "=" * 10)
    with torch.no_grad():
        total_samples = 0
        correct_samples = 0
        total_loss = 0
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)
            # data = data.reshape(data.shape[0], -1)
            result = model(data)
            loss = loss_fn(result, target)

            total_loss += loss.item()
            total_samples += target.shape[0]
            correct_samples += target.eq(result.argmax(-1)).sum().item()
        print(f"loss: {loss / total_samples}, accuracy: {correct_samples / total_samples * 100:.2f}% ")

# setup device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# hyperparameters
batch_size = 128
input_size = 784
num_classes = 10
in_channels = 1
learning_rate = 0.001
epochs = 10

# load data
train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
test_data = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=batch_size, shuffle=True)

# initialize network, loss, optimizer
model = CNN(in_channels, num_classes).to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

if __name__ == '__main__':
    train(model, loss_fn, optimizer, train_loader, epochs)
    test(model, loss_fn, test_loader)





